/*
 * Copyright (c) 2013-2018 上海汇数数据. All rights reserved.
 * @(#) BaseAsymmetricCipherApi.java 2018-08-06 17:26
 */

package cn.jh.common.core.crypto;


import cn.jh.common.core.base.Base64Kt;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.spec.AlgorithmParameterSpec;
import java.util.Enumeration;

import static com.google.common.base.Preconditions.checkArgument;


/**
 * 基于JCA体系的非对称加密算法的 `CipherApi` 接口的基础实现。
 *
 * @author Fuchun
 * @since 1.0
 */
public abstract class BaseAsymmetricCipherApi extends JcaCipherApi implements AsymmetricCipherApi {

    protected BaseAsymmetricCipherApi(AsymmetricCipherParams params) {
        super(params);
    }

    @Override
    public AsymmetricCipherParams getParams() {
        return (AsymmetricCipherParams) super.getParams();
    }

    @Override
    protected AlgorithmParameterSpec toParameterSpec(byte[] parameters) {
        return null;
    }

    public KeyPair generateNewKeyPair() {
        return generateNewKeyPair(getParams().getKeySize());
    }

    public KeyPair generateNewKeyPair(int keySize) {
        KeyPairGenerator kpg = KeyKt.getAndInitBcKpGen(getParams().getAlgorithm(), KeyBit.of(keySize));
        return kpg.generateKeyPair();
    }

//    @Override
//    protected byte[] encrypt(byte[] plaintext, Key key, @Nullable @SuppressWarnings("unused") byte[] parameters) {
//        checkArgument(key instanceof PublicKey, "The key must be a public key.");
//        return cryptWithKey(plaintext, key, parameters, Cipher.ENCRYPT_MODE);
//    }
//
//    @Override
//    protected byte[] encrypt(byte[] plaintext, byte[] key, @Nullable byte[] parameters) {
//        return super.encrypt(plaintext, key, parameters);
//    }
//
//    @Override
//    protected byte[] decrypt(byte[] cipherText, Key key, @Nullable @SuppressWarnings("unused") byte[] parameters) {
//        checkArgument(key instanceof PrivateKey, "The key must be a private key.");
//        return cryptWithKey(cipherText, key, parameters, Cipher.DECRYPT_MODE);
//    }

    @Override
    public byte[] sign(byte[] content, PrivateKey privateKey) throws CryptoException {
        checkArgument(content != null, "The content must be not null");
        checkArgument(privateKey != null, "The privateKey must be not null");
        try {
            Signature signature = newSignatureInstance();

            if (params.getSecureRandom() != null) {
                signature.initSign(privateKey, params.getSecureRandom());
            } else {
                signature.initSign(privateKey);
            }
            signature.initSign(privateKey);
            signature.update(content);

            return signature.sign();
        } catch (InvalidKeyException ex) {
            throw new CryptoException("Invalid " + params.getAlgorithm() + " private key(base64): \n" +
                    Base64Kt.encodeToMime(privateKey.getEncoded()), ex);
        } catch (SignatureException ex) {
            throw new CryptoException("Signature sign error: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()), ex);
        }
    }

    @Override
    public boolean verify(byte[] content, byte[] sign, PublicKey publicKey) {
        if (content == null || sign == null || publicKey == null) return false;
        Signature signature = newSignatureInstance();
        try {
            signature.initVerify(publicKey);
            signature.update(content);
            return signature.verify(sign);
        } catch (InvalidKeyException ex) {
            throw new CryptoException("Invalid " + params.getAlgorithm() + " public key(base64): \n" +
                    Base64Kt.encodeToMime(publicKey.getEncoded()), ex);
        } catch (SignatureException ex) {
            throw new CryptoException("Signature verify error: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()), ex);
        }
    }

    /**
     * 将指定的密钥数据转换为 {@code PublicKey}。
     *
     * @param key 要转换的密钥数据。
     * @return 指定密钥表示的 {@code PublicKey}。
     * @throws CryptoException          如果指定密钥格式错误。
     * @throws IllegalArgumentException 如果有 {@code key == null || key.length == 0}，或者
     * 当前非对称加密接口参数中指定的算法{@code params.algorithm} 错误。
     */
    protected PublicKey toPublicKey(byte[] key) throws CryptoException, IllegalArgumentException {
        return KeyKt.toPublicKey(key, params.getAlgorithm());
    }

    /**
     * 将指定的密钥数据转换为 {@code PrivateKey}。
     *
     * @param key 要转换的密钥数据。
     * @return 指定密钥表示的 {@code PrivateKey}。
     * @throws CryptoException          如果指定密钥格式错误。
     * @throws IllegalArgumentException 如果有 {@code key == null || key.length == 0}，或者
     * 当前非对称加密接口参数中指定的算法{@code params.algorithm} 错误。
     */
    protected PrivateKey toPrivateKey(byte[] key) throws CryptoException, IllegalArgumentException {
        return KeyKt.toPrivateKey(key, params.getAlgorithm());
    }

    public KeyPair loadKeyPair(@NotNull String keyPath, @NotNull String alias, @Nullable String keyPassed) {
        char[] password = keyToCharArray(keyPassed);
        try (FileInputStream fis = new FileInputStream(keyPath)) {
            KeyStore ks = loadKeyStore(fis, password);

            PrivateKey privateKey = (PrivateKey) ks.getKey(alias, password);
            Certificate certificate = ks.getCertificate(alias);
            PublicKey publicKey = certificate.getPublicKey();
            return new KeyPair(publicKey, privateKey);
        } catch (UnrecoverableKeyException ex) {
            throw new CryptoException("The key cannot be recovered in keyPath(\"" + keyPath +
                    "\"). Maybe password is wrong ?");
        } catch (Exception ex) {
            throw new CryptoException("Could not load private key from \"" + keyPath + "\". Exception: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()));
        }
    }

    /**
     * 加载指定证书或密钥文件中的私钥。
     *
     * @param keyPath   证书或密钥等文件路径。
     * @param keyPassed 证书或密钥的保护密码。
     * @return 返回证书或密钥文件中的私钥。
     * @throws CryptoException 如果读取私钥过程发生任何错误，包括但不限于
     *                         不支持 PKCS12，发生任何I/O错误，算法不能检验密钥库的完整性，证书不能被加载，保护密码错误等等。
     */
    public PrivateKey loadPrivateKey(@NotNull String keyPath, @Nullable String keyPassed) {
        char[] password = keyToCharArray(keyPassed);
        try (FileInputStream fis = new FileInputStream(keyPath)) {
            KeyStore ks = loadKeyStore(fis, password);

            String alias = findPrivateKeyAlias(ks);
            return (PrivateKey) ks.getKey(alias, password);
        } catch (UnrecoverableKeyException ex) {
            throw new CryptoException("The key cannot be recovered in keyPath(\"" + keyPath +
                    "\"). Maybe password is wrong ?");
        } catch (Exception ex) {
            throw new CryptoException("Could not load private key from \"" + keyPath + "\". Exception: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()));
        }
    }

    /**
     * 从输入流数据中加载指定证书或密钥的私钥。
     *
     * @param input     证书或密钥数据或文件流。
     * @param keyPassed 证书或密钥的保护密码。
     * @return 返回证书或密钥数据中的私钥。
     * @throws CryptoException 如果读取私钥过程发生任何错误，包括但不限于
     *                         不支持 PKCS12，发生任何I/O错误，算法不能检验密钥库的完整性，证书不能被加载，保护密码错误等等。
     */
    public PrivateKey loadPrivateKey(@NotNull InputStream input, @Nullable String keyPassed) {
        char[] password = keyToCharArray(keyPassed);
        KeyStore ks = loadKeyStore(input, password);
        try {
            String alias = findPrivateKeyAlias(ks);
            return (PrivateKey) ks.getKey(alias, password);
        } catch (UnrecoverableKeyException ex) {
            throw new CryptoException("The key cannot be recovered Maybe password is wrong ?");
        } catch (Exception ex) {
            throw new CryptoException("Could not load private key. Exception: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()));
        }
    }

    public PublicKey loadPublicKey(@NotNull String keyPath, @Nullable String keyPassed) {
        char[] password = keyToCharArray(keyPassed);
        try (FileInputStream fis = new FileInputStream(keyPath)) {
            KeyStore ks = loadKeyStore(fis, password);

            String alias = findPublicKeyAlias(ks);
            Certificate certificate = ks.getCertificate(alias);
            return certificate.getPublicKey();
        } catch (Exception ex) {
            throw new CryptoException("Could not load public key from \"" + keyPath + "\". Exception: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()));
        }
    }

    public PublicKey loadPublicKey(@NotNull InputStream input, @Nullable String keyPassed) {
        char[] password = keyToCharArray(keyPassed);
        KeyStore ks = loadKeyStore(input, password);
        try {
            String alias = findPublicKeyAlias(ks);
            Certificate certificate = ks.getCertificate(alias);
            return certificate.getPublicKey();
        } catch (Exception ex) {
            throw new CryptoException("Could not load public key. Exception: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()));
        }
    }

    public PublicKey loadKeyFromCertificate(@NotNull InputStream input) {
        try {
            CertificateFactory certificateFactory = CertificateFactory.getInstance("X.509");

            Certificate certificate = certificateFactory.generateCertificate(input);
            return certificate.getPublicKey();
        } catch (CertificateException ex) {
            throw new CryptoException("Could not load public key. Exception: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()));
        }
    }

    private char[] keyToCharArray(@Nullable String keyPassed) {
        char[] password;
        if (keyPassed == null || (keyPassed = keyPassed.trim()).isEmpty()) {
            password = null;
        } else {
            password = keyPassed.toCharArray();
        }
        return password;
    }

    private String findPublicKeyAlias(KeyStore ks) throws KeyStoreException {
        Enumeration<String> aliases = ks.aliases();
        String alias = null;
        while (aliases.hasMoreElements()) {
            alias = aliases.nextElement();
            if (ks.isCertificateEntry(alias)) {
                break;
            }
        }
        if (alias == null) {
            throw new CryptoException("The KeyStore could not load public key.");
        }
        return alias;
    }

    private String findPrivateKeyAlias(KeyStore ks) throws KeyStoreException {
        Enumeration<String> aliases = ks.aliases();
        String alias = null;
        while (aliases.hasMoreElements()) {
            alias = aliases.nextElement();
            if (ks.isKeyEntry(alias)) {
                break;
            }
        }
        if (alias == null) {
            throw new CryptoException("The KeyStore could not load private key.");
        }
        return alias;
    }

    private KeyStore loadKeyStore(@NotNull InputStream input, @Nullable char[] password) {
        try {
            KeyStore ks = KeyStore.getInstance("PKCS12");

            ks.load(input, password);

            return ks;
        } catch (KeyStoreException ex) {
            // 基本不会发生此异常
            throw new CryptoException("The KeyStore is invalid for `PKCS12`");
        } catch (IOException ex) {
            throw new CryptoException(ex);
        } catch (NoSuchAlgorithmException ex) {
            throw new CryptoException("The algorithm used to check the integrity of the keystore cannot be found.", ex);
        } catch (CertificateException ex) {
            throw new CryptoException("The keystore could not load any of certificates.", ex);
        } catch (Exception ex) {
            throw new CryptoException("Could not load KeyStore to use PKCS12. Exception: " +
                    (ex.getLocalizedMessage() != null ? ex.getLocalizedMessage() : ex.getMessage()), ex);
        }
    }

    private Signature newSignatureInstance() throws IllegalArgumentException {
        try {
            return Signature.getInstance(
                    getParams().getSignatureAlgorithm(), params.getProvider());
        } catch (NoSuchProviderException ex) {
            throw new IllegalArgumentException("No such Provider('" + params.getProvider() + "') on this platform.");
        } catch (NoSuchAlgorithmException ex) {
            throw new IllegalArgumentException("No such Algorithm('" + getParams().getSignatureAlgorithm() +
                    "') on this platform.", ex);
        }
    }
}
